import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D


def plotmaze(obstacle, startpoint, endpoint):  # 仅画地图
    fig = plt.figure()
    ax = plt.subplot(projection='3d')
    x = [d[0] for d in obstacle]
    y = [d[1] for d in obstacle]
    z = [d[2] for d in obstacle]

    ax.bar3d(x,
             y,
             np.zeros_like(z),
             dx=1,
             dy=1,
             dz=z,
             color='red')
    ax.bar3d([startpoint[0] - 0.5, endpoint[0] - 0.5],
             [startpoint[1] - 0.5, endpoint[1] - 0.5],
             [startpoint[2], endpoint[2]],
             dx=1,
             dy=1,
             dz=1,
             color='green')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.show()


def plotpoint(obstacle, startpoint, endpoint, pointlist):  # 画出路径图与地图（可用于训练时观测）
    if not pointlist:
        return
    fig = plt.figure()
    ax = plt.subplot(projection='3d')
    x = [d[0] for d in obstacle]
    y = [d[1] for d in obstacle]
    z = [d[2] for d in obstacle]

    ax.bar3d(x,
             y,
             np.zeros_like(z),
             dx=1,
             dy=1,
             dz=z,
             color='red',
             alpha=0.3)
    ax.bar3d([startpoint[0] - 0.5, endpoint[0] - 0.5],
             [startpoint[1] - 0.5, endpoint[1] - 0.5],
             [startpoint[2], endpoint[2]],
             dx=1,
             dy=1,
             dz=1,
             color='green')

    xpos = [c[0] for c in pointlist]
    ypos = [c[1] for c in pointlist]
    zpos = [c[2] for c in pointlist]
    ax.scatter(xpos,
               ypos,
               zpos,
               s=10,
               c='blue')

    ax.plot(xpos,
            ypos,
            zpos,
            c='black',
            label='path')

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.pause(1)
    plt.show(block=False)
    plt.close("all")


def pathpoint(obstacle, startpoint, endpoint, pointlist):  # 画出路径图与地图
    if not pointlist:
        print("none")
        return
    fig = plt.figure(2)
    ax = plt.subplot(projection='3d')
    x = [d[0] for d in obstacle]
    y = [d[1] for d in obstacle]
    z = [d[2] for d in obstacle]

    ax.bar3d(x,
             y,
             np.zeros_like(z),
             dx=1,
             dy=1,
             dz=z,
             color='red',
             alpha=0.2)
    ax.bar3d([startpoint[0] - 0.5, endpoint[0] - 0.5],
             [startpoint[1] - 0.5, endpoint[1] - 0.5],
             [startpoint[2], endpoint[2]],
             dx=1,
             dy=1,
             dz=1,
             color='green')

    xpos = [c[0] for c in pointlist]
    ypos = [c[1] for c in pointlist]
    zpos = [c[2] for c in pointlist]
    ax.scatter(xpos,
               ypos,
               zpos,
               s=10,
               c='blue')

    ax.plot(xpos,
            ypos,
            zpos,
            c='black',
            label='path')

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.show()
